import threading
from typing import Optional, Dict, Any
import queue
import numpy as np
import zarr
import time


class TrajectoryDataLogger:

    def __init__(self, path='demo_data.zarr', start_id=0, overwrite=False):
        self.root = zarr.open(path, mode='w') if overwrite else zarr.open(path, mode='a')
        self.q = queue.Queue()
        self.worker = threading.Thread(target=self._worker, daemon=True)
        self.worker.start()
        self.counter = start_id
        self.lock = threading.Lock()
        self.current_traj = []

    def start_trajectory(self, meta: Optional[dict] = None):
        self.current_traj = []
        self.current_meta = meta or {}

    def record_frame(self, 
                     color: np.ndarray, 
                     depth: np.ndarray, 
                     gripper_distance,
                     target_pose,
                     tactile_umi_left_src,
                     tactile_umi_left_force,
                     timestamp: Optional[float] = None):
        self.current_traj.append({
            'color': color,
            'depth': depth,
            'gripper_distance':gripper_distance,
            "target_pose":target_pose,
            'tactile_umi_left_src': tactile_umi_left_src,
            'tactile_umi_left_force': tactile_umi_left_force,
            'timestamp': timestamp if timestamp else time.time()
        })

    def end_trajectory(self):
        with self.lock:
            self.counter += 1
            traj_id = f"trajectory_{self.counter:06d}"
        self.q.put((traj_id, self.current_traj, self.current_meta))
        self.current_traj = []

    def _worker(self):
        while True:
            task = self.q.get()
            if task is None:
                break
            traj_id, traj_data, meta = task
            print(f"Save trajectory {traj_id}, do not exit!!!")
            g = self.root.require_group(traj_id)

            color_seq = np.stack([f['color'] for f in traj_data])
            depth_seq = np.stack([f['depth'] for f in traj_data])
            gripper_dist_seq = np.stack([f['gripper_distance'] for f in traj_data])
            target_pose_seq = np.stack([f['target_pose'] for f in traj_data])
            rgb_seq = np.stack([f['tactile_umi_left_src'] for f in traj_data])
            force_seq = np.stack([f['tactile_umi_left_force'] for f in traj_data])
            timestamps = np.array([f['timestamp'] for f in traj_data])

            g.create_dataset('depth', data=depth_seq, chunks=True, overwrite=True)
            g.create_dataset('color', data=color_seq, chunks=True, overwrite=True)
            g.create_dataset('target_pose', data=target_pose_seq, chunks=True, overwrite=True)
            g.create_dataset('gripper_distance', data=gripper_dist_seq, chunks=True, overwrite=True)
            g.create_dataset('tactile_umi_left_src', data=rgb_seq, chunks=True, overwrite=True)
            g.create_dataset('tactile_umi_left_force', data=force_seq, chunks=True, overwrite=True)
            g.create_dataset('timestamp', data=timestamps, chunks=True, overwrite=True)
            g.attrs.update(meta)
            print(f"data saved to {traj_id}.zarr, save exit.")

    def close(self):
        self.q.put(None)
        self.worker.join()
